from utils.polygon_judge import *
import copy
import numpy as np

from config.config_setting import rect_size
from utils import img_util

"""
越界检测
由于小组中已经有成型的越界检测项目代码，聚集人群识别用类似的思路。
1 在视频中画出多边形
2 利用已有模型检测出人框
3 判断人框与多边形的边缘是否相交并计数。
4 如果第3步中计数的数量大于预定的阈值，则将画面中人框的颜色标为红色并提示
第97行，可写死tpPointsChoose的值，用来固定多边形的区域。
"""

algo_id = 40


def draw_box_by_data_trt(boxes, scores, label_id, label_dict, img, param, device):
    #print(type(boxes), type(scores), type(label_id))
    #print(boxes, scores, label_id)
    y_len, x_len = img.shape[:2]
    #print(y_len, x_len)
    person_num, polypoint = None, []
    #print(param)
    #print(list(param.keys()))
    if "rule" in list(param.keys()):
        person_num = param["rule"]["personNum"]
        polygonPoint = param["rule"]["polygonPoint"]
        for i in range(len(polygonPoint)):
            point = (int(polygonPoint[i]["x"]*x_len/rect_size[0]), int(polygonPoint[i]["y"]*y_len/rect_size[1]))
            polypoint.append(point)

        #print(person_num,"+", polypoint)

    cv2.polylines(img, np.int32([polypoint]), True, (0,0,255), thickness=3)


    save = False
    crowed_num = 0  # 与多边形相交的矩形个数，即聚集人群数
    #boxes, scores, label_id = np.squeeze(boxes), np.squeeze(scores), np.squeeze(label_id).astype(np.int32)

    for (box, score, idx) in zip(boxes.tolist(), scores.tolist(), label_id.tolist()):
        if idx in label_dict:
            label = label_dict[idx]
            if score > 0.65 and label == 'person':
                #print("change before:",box[0], box[1], box[2], box[3])
                x_min, y_min, x_max, y_max = int(box[0]), int(box[1]), int(box[2]), int(box[3])
                #print("after change:", x_min, y_min, x_max, y_max)
                if polypoint and person_num and intersection_rect_poly(polypoint, (x_max, x_min), (y_max, y_min)):
                    #print("111")
                    crowed_num += 1
                    if crowed_num >= person_num:
                        save = True
                        cv2.rectangle(img, (x_min, y_min), (x_max, y_max), (0, 0, 255), thickness=3)
                    else:
                        cv2.rectangle(img, (x_min, y_min), (x_max, y_max), (0, 255, 0), thickness=3)
                else:
                    #print("222")
                    cv2.polylines(img, np.int32([polypoint]), True, (0,0,255), thickness=3)


    cv2.waitKey(1)
    # 对识别结果做去重操作
    if save:
        ori_img = copy.deepcopy(img)
        img_util.save_image(img, device, algo_id, ori_img)

    return img


if __name__ == '__main__':
    pass
